package com.gitee.zycra.file.utils;

import com.gitee.zycra.file.param.DownloadParam;
import com.gitee.zycra.file.param.UploadParam;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;

/**
 * Util for encrypt or decrypt the file to transfer between local and SFTP server.
 * <p>Encryption scheme:
 * <p>1.Generate an AES key for encrypt the file.
 * <p>2.Encrypt the file with AES by generated AES key.
 * <p>3.Zstd compress the encrypted file.
 * <p>4.Encrypt the generated AES key with outer SM2 public key.
 * <p>5.Sign the generated AES key with inner SM2 private key.
 * <p>6.Add the concatenated string of encrypted AES key and sign result at the end of compressed file.
 * <p>Concatenate scheme: zst file content+encrypted key+key sign result+2 byte encrypted key length+2 byte sign result length.
 *
 * @author zycra
 * @since 1.0.0
 */
public final class FileEncryptUtil {

    private static final Logger LOGGER = LoggerFactory.getLogger(FileEncryptUtil.class);

    /**
     * Temp AES encrypted file extention.
     *
     * @since 1.0.0
     */
    private static final String AES_ENCRYPT_SUFFIX = ".aes";

    /**
     * Temp zstd compressed file extention.
     *
     * @since 1.0.0
     */
    private static final String ZSTD_SUFFIX = ".zst";

    /**
     * Read and write auth.
     *
     * @since 1.0.0
     */
    private static final String RW = "rw";

    /**
     * Tail buffer size length for seek.
     *
     * @since 1.0.0
     */
    private static final int EXTEND_BUFFER_SIZE_LENGTH = 4;

    private FileEncryptUtil() {
    }

    /**
     * Encrypt the file before upload.
     *
     * @param uploadParam upload param.
     * @return file path of the encrypted file.
     * @since 1.0.0
     */
    public static String encrypt(UploadParam uploadParam) {
        if (!uploadParam.isEncrypt()) {
            return uploadParam.getLocalFilePathAndName();
        }
        String localFilePathAndName = uploadParam.getLocalFilePathAndName();
        String aesEncryptFilePathAndName = localFilePathAndName + AES_ENCRYPT_SUFFIX;
        String zstdFilePathAndName = localFilePathAndName + ZSTD_SUFFIX;
        File originFile = new File(localFilePathAndName);
        File aesEncryptFile = new File(aesEncryptFilePathAndName);
        File zstdFile = new File(zstdFilePathAndName);
        String aesKey = AESUtil.generateKey();
        boolean encryptSuccess = AESUtil.encrypt(originFile, aesEncryptFile, aesKey);
        if (!encryptSuccess) {
            LOGGER.error("AES encrypt failed, localFilePathAndName={}", localFilePathAndName);
            return null;
        }
        if (uploadParam.isDeleteOriginFile()) {
            try {
                Files.delete(originFile.toPath());
            } catch (IOException e) {
                LOGGER.error("origin file delete error, localFilePathAndName={}", localFilePathAndName, e);
            }
        }
        if (!CompressUtil.zstdCompress(aesEncryptFile, zstdFile)) {
            LOGGER.error("zstd compress file failed, localFilePathAndName={}", localFilePathAndName);
            return null;
        }
        try {
            Files.delete(aesEncryptFile.toPath());
        } catch (IOException e) {
            LOGGER.error("AES encrypt file delete error, aesEncryptFilePathAndName={}", aesEncryptFilePathAndName, e);
        }
        String encryptKey = SM2Util.encryptByPublicKey(aesKey, uploadParam.getEncryptKey());
        String sign = SM2Util.signByPrivateKey(aesKey, uploadParam.getSignKey());
        if (encryptKey == null || sign == null) {
            LOGGER.error("SM2 encrypt or sign error");
            return null;
        }
        try (RandomAccessFile raf = new RandomAccessFile(zstdFile, RW)) {
            raf.seek(raf.length());
            raf.write(encryptKey.getBytes(StandardCharsets.UTF_8));
            raf.write(sign.getBytes(StandardCharsets.UTF_8));
            raf.writeShort(encryptKey.length());
            raf.writeShort(sign.length());
        } catch (Exception e) {
            LOGGER.error("write encrypt key and sign error", e);
        }
        return zstdFilePathAndName;
    }

    /**
     * Decrypt the file after download.
     *
     * @param downloadParam download param.
     * @param localFileName origin file path to decrypt.
     * @return whether successful.
     * @since 1.0.0
     */
    public static boolean decrypt(DownloadParam downloadParam, String localFileName) {
        if (!downloadParam.isDecrypt()) {
            return true;
        }
        File originFile = new File(localFileName);
        String aesKey = null;
        try (RandomAccessFile raf = new RandomAccessFile(originFile, RW)) {
            raf.seek(raf.length() - EXTEND_BUFFER_SIZE_LENGTH);
            short encryptLength = raf.readShort();
            short signLength = raf.readShort();
            long cutFileLength = raf.length() - EXTEND_BUFFER_SIZE_LENGTH - signLength - encryptLength;
            raf.seek(cutFileLength);
            byte[] encryptArr = new byte[encryptLength];
            byte[] signArr = new byte[signLength];
            raf.read(encryptArr);
            raf.read(signArr);
            String encryptKey = new String(encryptArr, StandardCharsets.UTF_8);
            String sign = new String(signArr, StandardCharsets.UTF_8);
            aesKey = SM2Util.decryptByPrivateKey(encryptKey, downloadParam.getDecryptKey());
            boolean checkSuccess = SM2Util.verifyByPublicKey(aesKey, downloadParam.getCheckSignKey(), sign);
            if (!checkSuccess) {
                LOGGER.error("file check sign failed, localFileName={}", localFileName);
                return false;
            }
            raf.seek(0L);
            raf.setLength(cutFileLength);
        } catch (Exception e) {
            LOGGER.error("decrypt file error", e);
            return false;
        }
        File aesEncryptFile = new File(localFileName + AES_ENCRYPT_SUFFIX);
        if (!CompressUtil.zstdUnCompress(originFile, aesEncryptFile)) {
            LOGGER.error("zstd uncompress file failed, localFileName={}", localFileName);
            return false;
        }
        try {
            Files.delete(originFile.toPath());
        } catch (IOException e) {
            LOGGER.error("delete zstd file error", e);
        }
        boolean result = AESUtil.decrypt(aesEncryptFile, originFile, aesKey);
        try {
            Files.delete(aesEncryptFile.toPath());
        } catch (IOException e) {
            LOGGER.error("delete aes encrypt file error", e);
        }
        return result;
    }
}
